--[[
  * Description: A Weimob Simple WAF
  * Author : Luma
  * Email : 825244020@qq.com
  * Version: 0.1
  * Date : 2015/09/18
--]]
local iputils  = require "iputils" -- iputils another lua module

iputils.enable_lrucache()

local WeimobWAF = {}

function WeimobWAF:new(config)
    local ob = config
    if ob == nil then
        ob = {
            logfile = "/tmp/xx.log"
        }
    end
    self.__index = self
    return setmetatable(ob, self)
end

function WeimobWAF:getKeys(tab) 
    local keys={}
    local n=0
    if type(tab) == "table" then
       for k,v in pairs(tab) do
          n=n+1
          keys[n]=k
       end
       return keys
    else
       return nil
    end 
end

function WeimobWAF:hasKey(t,k)
    if type(t) == "table" then
        local keys = self:getKeys(t)
        for i=1,#keys do 
            if keys[i] == k then return true end
        end
        return false
    end
end

function WeimobWAF:checkUA(index,t)  
    local u = ngx.var.http_user_agent
    local m = ngx.re.match(u,t["ua"])
    if m then
        return true
    else
        return false
    end

end

function WeimobWAF:checkIP(index,t)
    local ip = self:getRealIp()
    local ipBlockList = iputils.parse_cidrs(self:explode(",",t["ip"])) 
    if iputils.ip_in_cidrs(ip,ipBlockList) then 
       return true
    else
       return false
    end 
end

function WeimobWAF:checkHost(index,t)
    local host = ngx.var.http_host
    if t["host"] == host then
        return true
    else
        return false
    end
end

function WeimobWAF:checkArgs(index,t)
    local method=ngx.req.get_method()
    if method ~= "GET" then
        return false
    else
        local args = ngx.req.get_uri_args()
        if args ~= nil then 
            local point,n = {},1
            for k,v in string.gmatch(t["args"], "(%w+)=(%w+)") do
                if v ~= args[k] then 
                    point[n] = false
                else 
                    point[n] = true
                end 
                n = n+1
            end
            for _,v in pairs(point) do
               if v ~= true then return false  end
            end
            return true
        else
            return false
        end    
    end
end

function WeimobWAF:processRules(index,t)
    local m = {
        ua = "checkUA",
        ip = "checkIP",
        host = "checkHost",
        args = "checkArgs",
    }
    local u = {}
    for k,v in ipairs(self:getKeys(t)) do
       u[k]=self[m[v]](self,index,t)
    end 
    for _,v in pairs(u) do
        if v ~= true then return false  end
    end 
    return true 
end


function WeimobWAF:explode(sep,str)
    if (sep=='') then return false end
    local pos,arr = 0,{}
    for st,sp in function() return string.find(str,sep,pos,true) end do
        table.insert(arr,string.sub(str,pos,st-1))
        pos = sp + 1
    end
    table.insert(arr,string.sub(str,pos))
    return arr
end



function WeimobWAF:getRealIp(field)
    local h = ngx.req.get_headers()
    if field == nil then
        if h["x-forwarded-for"] then
            for k,v in self:explode(",",h["x-forwarded-for"]) do
                if k == 1 then return v  end
            end
        elseif h["x-real-ip"] then
            return h["x-real-ip"]
        else
            return ngx.var.remote_addr
        end
    else
        for k,v in pairs(h) do
            if k == field  then return v end
        end
    end
end

function WeimobWAF:writeLog(logfile,msg)
    local file =  logfile or self.logfile
    local fd = io.open(file,"ab")
    fd:write(msg)
    fd:flush()
    fd:close()

end

function WeimobWAF:Log(method,ridx)
    local ip = self:getRealIp()
    local ua = ngx.var.http_user_agent
    local url = ngx.var.server_name..ngx.var.request_uri
    local time = ngx.localtime()
    local line = ip.."  "..time.." "..method.." "..url.."  "..ua.."  ".."[".."rule:"..ridx.."  matched]".."\n"
    self:writeLog(self.logfile,line)
end


function WeimobWAF:checkRules(rules)
    if  rules == nil then
        return false
    else
        for i=1,#rules do
            local m = rules[i] --why???
            local rt = self:processRules(i,m)
            if not rt then 
                goto continue 
            else 
                self:Log("GET",i,"xxx") 
                return true 
            end
            ::continue:: 
        end
        return false
    end
end

--local ins = C:new{logfile:"/temp/waf.log"}
local ins = WeimobWAF:new()
return ins
